package heap

type Heap[T any] struct {
	data       []T
	bigHeap    bool
	comparator func(v1, v2 T) bool // 小于true,其他false
	zeroValue  T
}

func NewHeap[T any](bigHeap bool, c func(v1, v2 T) bool) *Heap[T] {
	return &Heap[T]{
		data:       make([]T, 0, 10),
		bigHeap:    bigHeap,
		comparator: c,
	}
}

func (h *Heap[T]) parent(i int) int {
	if i < 0 {
		return i
	}
	return i >> 1
}

func (h *Heap[T]) Push(val T) {
	h.data = append(h.data, val)
	h.topDown()
}

func (h *Heap[T]) Pop() (T, bool) {
	if len(h.data) == 0 {
		return h.zeroValue, false
	}
	top := h.data[0]
	h.data[0] = h.data[len(h.data)-1]
	h.data = h.data[:len(h.data)-1]
	h.downTop(0)
	return top, true
}

func (h *Heap[T]) Empty() bool {
	return len(h.data) == 0
}

func (h *Heap[T]) topDown() {
	for i := range h.data {
		current := i
		parent := h.parent(current - 1)

		for parent >= 0 && !h.compare(h.data[current], h.data[parent]) {
			h.data[current], h.data[parent] = h.data[parent], h.data[current]
			current = parent
			parent = h.parent(current - 1)
		}
	}
}

func (h *Heap[T]) downTop(index int) {
	n := len(h.data) - 1
	for {
		suiteIndex := index
		left := index<<1 + 1
		right := left + 1

		if left <= n && !h.compare(h.data[left], h.data[suiteIndex]) {
			suiteIndex = left
		}

		if right <= n && !h.compare(h.data[right], h.data[suiteIndex]) {
			suiteIndex = right
		}

		if suiteIndex == index {
			break
		}
		h.data[index], h.data[suiteIndex] = h.data[suiteIndex], h.data[index]
		index = suiteIndex
	}
}

func (h *Heap[T]) compare(v1, v2 T) bool {
	if h.bigHeap {
		// 大根堆，当前节点小于父节点即可
		return h.comparator(v1, v2)
	}
	return !h.comparator(v1, v2)
}
